package batch

import (
	"context"

	"github.com/Jeffail/benthos/v3/lib/message"
	"github.com/Jeffail/benthos/v3/lib/types"
)

type batchedCountKeyType int

const batchedCountKey batchedCountKeyType = iota

// CollapsedCount attempts to extract the actual number of messages that were
// collapsed into the resulting message part. This value could be greater than 1
// when users configure processors that archive batched message parts.
func CollapsedCount(p types.Part) int {
	if v, ok := message.GetContext(p).Value(batchedCountKey).(int); ok {
		return v
	}
	return 1
}

// MessageCollapsedCount attempts to extract the actual number of messages that
// were combined into the resulting batched message parts. This value could
// differ from message.Len() when users configure processors that archive
// batched message parts.
func MessageCollapsedCount(m types.Message) int {
	total := 0
	m.Iter(func(i int, p types.Part) error {
		total += CollapsedCount(p)
		return nil
	})
	return total
}

// WithCollapsedCount returns a message part with a context indicating that this
// message is the result of collapsing a number of messages. This allows
// downstream components to know how many total messages were combined.
func WithCollapsedCount(p types.Part, count int) types.Part {
	// Start with the previous length which could also be >1.
	ctx := message.GetContext(p)
	base := 1
	if v, ok := ctx.Value(batchedCountKey).(int); ok {
		base = v
	}

	// The new length is the previous length plus the total messages put into
	// this batch (minus one to prevent double counting the original part).
	ctx = context.WithValue(ctx, batchedCountKey, base-1+count)

	return message.WithContext(ctx, p)
}
